/* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 * Slabs memory allocation, based on powers-of-N. Slabs are up to 1MB in size
 * and are divided into chunks. The chunk sizes start off at the size of the
 * "item" structure plus space for a small key and value. They increase by
 * a multiplier factor from there, up to half the maximum slab size. The last
 * slab size is always 1MB, since that's the maximum item size allowed by the
 * memcached protocol.
 */
#include "../lib/memcached.h"
#include "../lib/hash.h"
#include <sys/stat.h>
#include <sys/signal.h>
#include <sys/resource.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <errno.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <assert.h>
#include <pthread.h>

/* powers-of-N allocation structures */

typedef struct
{
    unsigned int size; /* sizes of items */
    unsigned int perslab; /* how many items per slab */

    void *slots; /* list of item ptrs */
    unsigned int sl_curr; /* total free items in list */

    unsigned int slabs; /* how many slabs were allocated for this class */

    void **slab_list; /* array of slab pointers */
    unsigned int list_size; /* size of prev array */

    unsigned int killing; /* index+1 of dying slab, or zero if none */
    size_t requested; /* The number of requested bytes */
} slabclass_t;

static slabclass_t slabclass[MAX_NUMBER_OF_SLAB_CLASSES];
static size_t mem_limit = 0;
static size_t mem_malloced = 0;
/* If the memory limit has been hit once. Used as a hint to decide when to
 * early-wake the LRU maintenance thread */
static bool mem_limit_reached = false;
static int power_largest;

static void *mem_base = NULL;
static void *mem_current = NULL;
static size_t mem_avail = 0;

/**
 * Access to the slab allocator is protected by this lock
 */
static pthread_mutex_t slabs_lock = PTHREAD_MUTEX_INITIALIZER;
static pthread_mutex_t slabs_rebalance_lock = PTHREAD_MUTEX_INITIALIZER;

/*
 * Forward Declarations
 */
static int do_slabs_newslab ( const unsigned int id );
static void *memory_allocate ( size_t size );
static void do_slabs_free ( void *ptr, const size_t size, unsigned int id );

/* Preallocate as many slab pages as possible (called from slabs_init)
   on start-up, so users don't get confused out-of-memory errors when
   they do have free (in-slab) space, but no space to make new slabs.
   if maxslabs is 18 (POWER_LARGEST - POWER_SMALLEST + 1), then all
   slab types can be made.  if max memory is less than 18 MB, only the
   smaller ones will be made.  */
static void slabs_preallocate ( const unsigned int maxslabs );

/*
 * Figures out which slab class (chunk size) is required to store an item of
 * a given size.
 *
 * Given object size, return id to use when allocating/freeing memory for object
 * 0 means error: can't store such a large object
 */

unsigned int slabs_clsid ( const size_t size )
{
    int res = POWER_SMALLEST;

    if ( size == 0 )
        return 0;
    while ( size > slabclass[res].size )
        if ( res ++ == power_largest ) /* won't fit in the biggest slab */
            return 0;
    return res;
}

/**
 * Determines the chunk sizes and initializes the slab class descriptors
 * accordingly.
 */
void slabs_init ( const size_t limit, const double factor, const bool prealloc )
{
    int i = POWER_SMALLEST - 1;
    unsigned int size = sizeof (item ) + settings.chunk_size;

    mem_limit = limit;

    if ( prealloc )
    {
        /* Allocate everything in a big chunk with malloc */
        mem_base = malloc (mem_limit);
        if ( mem_base != NULL )
        {
            mem_current = mem_base;
            mem_avail = mem_limit;
        }
        else
        {
            fprintf (stderr, "Warning: Failed to allocate requested memory in"
                     " one large chunk.\nWill allocate in smaller chunks\n");
        }
    }

    memset (slabclass, 0, sizeof (slabclass ));

    while ( ++ i < MAX_NUMBER_OF_SLAB_CLASSES - 1 && size <= settings.item_size_max / factor )
    {
        /* Make sure items are always n-byte aligned */
        if ( size % CHUNK_ALIGN_BYTES )
            size += CHUNK_ALIGN_BYTES - ( size % CHUNK_ALIGN_BYTES );

        slabclass[i].size = size;
        slabclass[i].perslab = settings.item_size_max / slabclass[i].size;
        size *= factor;
        if ( settings.verbose > 1 )
        {
            fprintf (stderr, "slab class %3d: chunk size %9u perslab %7u\n",
                     i, slabclass[i].size, slabclass[i].perslab);
        }
    }

    power_largest = i;
    slabclass[power_largest].size = settings.item_size_max;
    slabclass[power_largest].perslab = 1;
    if ( settings.verbose > 1 )
    {
        fprintf (stderr, "slab class %3d: chunk size %9u perslab %7u\n",
                 i, slabclass[i].size, slabclass[i].perslab);
    }

    /* for the test suite:  faking of how much we've already malloc'd */
    {
        char *t_initial_malloc = getenv ("T_MEMD_INITIAL_MALLOC");
        if ( t_initial_malloc )
        {
            mem_malloced = ( size_t ) atol (t_initial_malloc);
        }

    }

    if ( prealloc )
    {
        slabs_preallocate (power_largest);
    }
}

static void slabs_preallocate ( const unsigned int maxslabs )
{
    int i;
    unsigned int prealloc = 0;

    /* pre-allocate a 1MB slab in every size class so people don't get
       confused by non-intuitive "SERVER_ERROR out of memory"
       messages.  this is the most common question on the mailing
       list.  if you really don't want this, you can rebuild without
       these three lines.  */

    for ( i = POWER_SMALLEST; i < MAX_NUMBER_OF_SLAB_CLASSES; i ++ )
    {
        if ( ++ prealloc > maxslabs )
            return;
        if ( do_slabs_newslab (i) == 0 )
        {
            fprintf (stderr, "Error while preallocating slab memory!\n"
                     "If using -L or other prealloc options, max memory must be "
                     "at least %d megabytes.\n", power_largest);
            exit (1);
        }
    }

}

static int grow_slab_list ( const unsigned int id )
{
    slabclass_t *p = & slabclass[id];
    if ( p->slabs == p->list_size )
    {
        size_t new_size = ( p->list_size != 0 ) ? p->list_size * 2 : 16;
        void *new_list = realloc (p->slab_list, new_size * sizeof (void * ));
        if ( new_list == 0 ) return 0;
        p->list_size = new_size;
        p->slab_list = new_list;
    }
    return 1;
}

static void split_slab_page_into_freelist ( char *ptr, const unsigned int id )
{
    slabclass_t *p = & slabclass[id];
    int x;
    for ( x = 0; x < p->perslab; x ++ )
    {
        do_slabs_free (ptr, 0, id);
        ptr += p->size;
    }
}

static int do_slabs_newslab ( const unsigned int id )
{
    slabclass_t *p = & slabclass[id];
    int len = settings.slab_reassign ? settings.item_size_max
        : p->size * p->perslab;
    char *ptr;

    if ( ( mem_limit && mem_malloced + len > mem_limit && p->slabs > 0 ) )
    {
        mem_limit_reached = true;
        MEMCACHED_SLABS_SLABCLASS_ALLOCATE_FAILED (id);
        return 0;
    }

    if ( ( grow_slab_list (id) == 0 ) ||
         ( ( ptr = memory_allocate (( size_t ) len) ) == 0 ) )
    {

        MEMCACHED_SLABS_SLABCLASS_ALLOCATE_FAILED (id);
        return 0;
    }

    memset (ptr, 0, ( size_t ) len);
    split_slab_page_into_freelist (ptr, id);

    p->slab_list[p->slabs ++] = ptr;
    mem_malloced += len;
    MEMCACHED_SLABS_SLABCLASS_ALLOCATE (id);

    return 1;
}

/*@null@*/
static void *do_slabs_alloc ( const size_t size, unsigned int id, unsigned int *total_chunks )
{
    slabclass_t *p;
    void *ret = NULL;
    item *it = NULL;

    if ( id < POWER_SMALLEST || id > power_largest )
    {
        MEMCACHED_SLABS_ALLOCATE_FAILED (size, 0);
        return NULL;
    }
    p = & slabclass[id];
    assert (p->sl_curr == 0 || ( ( item * ) p->slots )->slabs_clsid == 0);

    *total_chunks = p->slabs * p->perslab;
    /* fail unless we have space at the end of a recently allocated page,
       we have something on our freelist, or we could allocate a new page */
    if ( ! ( p->sl_curr != 0 || do_slabs_newslab (id) != 0 ) )
    {
        /* We don't have more memory available */
        ret = NULL;
    }
    else if ( p->sl_curr != 0 )
    {
        /* return off our freelist */
        it = ( item * ) p->slots;
        p->slots = it->next;
        if ( it->next ) it->next->prev = 0;
        /* Kill flag and initialize refcount here for lock safety in slab
         * mover's freeness detection. */
        it->it_flags &= ~ ITEM_SLABBED;
        it->refcount = 1;
        p->sl_curr --;
        ret = ( void * ) it;
    }

    if ( ret )
    {
        p->requested += size;
        MEMCACHED_SLABS_ALLOCATE (size, id, p->size, ret);
    }
    else
    {
        MEMCACHED_SLABS_ALLOCATE_FAILED (size, id);
    }

    return ret;
}

static void do_slabs_free ( void *ptr, const size_t size, unsigned int id )
{
    slabclass_t *p;
    item *it;

    assert (id >= POWER_SMALLEST && id <= power_largest);
    if ( id < POWER_SMALLEST || id > power_largest )
        return;

    MEMCACHED_SLABS_FREE (size, id, ptr);
    p = & slabclass[id];

    it = ( item * ) ptr;
    it->it_flags |= ITEM_SLABBED;
    it->slabs_clsid = 0;
    it->prev = 0;
    it->next = p->slots;
    if ( it->next ) it->next->prev = it;
    p->slots = it;

    p->sl_curr ++;
    p->requested -= size;
    return;
}

static int nz_strcmp ( int nzlength, const char *nz, const char *z )
{
    int zlength = strlen (z);
    return (zlength == nzlength ) && ( strncmp (nz, z, zlength) == 0 ) ? 0 : - 1;
}

static void *memory_allocate ( size_t size )
{
    void *ret;

    if ( mem_base == NULL )
    {
        /* We are not using a preallocated large memory chunk */
        ret = malloc (size);
    }
    else
    {
        ret = mem_current;

        if ( size > mem_avail )
        {
            return NULL;
        }

        /* mem_current pointer _must_ be aligned!!! */
        if ( size % CHUNK_ALIGN_BYTES )
        {
            size += CHUNK_ALIGN_BYTES - ( size % CHUNK_ALIGN_BYTES );
        }

        mem_current = ( ( char* ) mem_current ) + size;
        if ( size < mem_avail )
        {
            mem_avail -= size;
        }
        else
        {
            mem_avail = 0;
        }
    }

    return ret;
}

void *slabs_alloc ( size_t size, unsigned int id, unsigned int *total_chunks )
{
    void *ret;

    pthread_mutex_lock (&slabs_lock);
    ret = do_slabs_alloc (size, id, total_chunks);
    pthread_mutex_unlock (&slabs_lock);
    return ret;
}

void slabs_free ( void *ptr, size_t size, unsigned int id )
{
    pthread_mutex_lock (&slabs_lock);
    do_slabs_free (ptr, size, id);
    pthread_mutex_unlock (&slabs_lock);
}

void slabs_adjust_mem_requested ( unsigned int id, size_t old, size_t ntotal )
{
    pthread_mutex_lock (&slabs_lock);
    slabclass_t *p;
    if ( id < POWER_SMALLEST || id > power_largest )
    {
        fprintf (stderr, "Internal error! Invalid slab class\n");
        abort ();
    }

    p = & slabclass[id];
    p->requested = p->requested - old + ntotal;
    pthread_mutex_unlock (&slabs_lock);
}

unsigned int slabs_available_chunks ( const unsigned int id, bool *mem_flag,
                                      unsigned int *total_chunks )
{
    unsigned int ret;
    slabclass_t *p;

    pthread_mutex_lock (&slabs_lock);
    p = & slabclass[id];
    ret = p->sl_curr;
    if ( mem_flag != NULL )
        *mem_flag = mem_limit_reached;
    if ( total_chunks != NULL )
        *total_chunks = p->slabs * p->perslab;
    pthread_mutex_unlock (&slabs_lock);
    return ret;
}

static pthread_cond_t slab_rebalance_cond = PTHREAD_COND_INITIALIZER;
static volatile int do_run_slab_thread = 1;
static volatile int do_run_slab_rebalance_thread = 1;

#define DEFAULT_SLAB_BULK_CHECK 1
int slab_bulk_check = DEFAULT_SLAB_BULK_CHECK;

static int slab_rebalance_start ( void )
{
    slabclass_t *s_cls;
    int no_go = 0;

    pthread_mutex_lock (&slabs_lock);

    if ( slab_rebal.s_clsid < POWER_SMALLEST ||
         slab_rebal.s_clsid > power_largest ||
         slab_rebal.d_clsid < POWER_SMALLEST ||
         slab_rebal.d_clsid > power_largest ||
         slab_rebal.s_clsid == slab_rebal.d_clsid )
        no_go = - 2;

    s_cls = & slabclass[slab_rebal.s_clsid];

    if ( ! grow_slab_list (slab_rebal.d_clsid) )
    {
        no_go = - 1;
    }

    if ( s_cls->slabs < 2 )
        no_go = - 3;

    if ( no_go != 0 )
    {
        pthread_mutex_unlock (&slabs_lock);
        return no_go; /* Should use a wrapper function... */
    }

    s_cls->killing = 1;

    slab_rebal.slab_start = s_cls->slab_list[s_cls->killing - 1];
    slab_rebal.slab_end = ( char * ) slab_rebal.slab_start +
        ( s_cls->size * s_cls->perslab );
    slab_rebal.slab_pos = slab_rebal.slab_start;
    slab_rebal.done = 0;

    /* Also tells do_item_get to search for items in this slab */
    slab_rebalance_signal = 2;

    if ( settings.verbose > 1 )
    {
        fprintf (stderr, "Started a slab rebalance\n");
    }

    pthread_mutex_unlock (&slabs_lock);

    return 0;
}

enum move_status
{
    MOVE_PASS = 0, MOVE_FROM_SLAB, MOVE_FROM_LRU, MOVE_BUSY, MOVE_LOCKED
};

/* refcount == 0 is safe since nobody can incr while item_lock is held.
 * refcount != 0 is impossible since flags/etc can be modified in other
 * threads. instead, note we found a busy one and bail. logic in do_item_get
 * will prevent busy items from continuing to be busy
 * NOTE: This is checking it_flags outside of an item lock. I believe this
 * works since it_flags is 8 bits, and we're only ever comparing a single bit
 * regardless. ITEM_SLABBED bit will always be correct since we're holding the
 * lock which modifies that bit. ITEM_LINKED won't exist if we're between an
 * item having ITEM_SLABBED removed, and the key hasn't been added to the item
 * yet. The memory barrier from the slabs lock should order the key write and the
 * flags to the item?
 * If ITEM_LINKED did exist and was just removed, but we still see it, that's
 * still safe since it will have a valid key, which we then lock, and then
 * recheck everything.
 * This may not be safe on all platforms; If not, slabs_alloc() will need to
 * seed the item key while holding slabs_lock.
 */
static int slab_rebalance_move ( void )
{
    slabclass_t *s_cls;
    int x;
    int was_busy = 0;
    int refcount = 0;
    uint32_t hv;
    void *hold_lock;
    enum move_status status = MOVE_PASS;

    pthread_mutex_lock (&slabs_lock);

    s_cls = & slabclass[slab_rebal.s_clsid];

    for ( x = 0; x < slab_bulk_check; x ++ )
    {
        hv = 0;
        hold_lock = NULL;
        item *it = slab_rebal.slab_pos;
        status = MOVE_PASS;
        if ( it->slabs_clsid != 255 )
        {
            /* ITEM_SLABBED can only be added/removed under the slabs_lock */
            if ( it->it_flags & ITEM_SLABBED )
            {
                /* remove from slab freelist */
                if ( s_cls->slots == it )
                {
                    s_cls->slots = it->next;
                }
                if ( it->next ) it->next->prev = it->prev;
                if ( it->prev ) it->prev->next = it->next;
                s_cls->sl_curr --;
                status = MOVE_FROM_SLAB;
            }
            else if ( ( it->it_flags & ITEM_LINKED ) != 0 )
            {
                /* If it doesn't have ITEM_SLABBED, the item could be in any
                 * state on its way to being freed or written to. If no
                 * ITEM_SLABBED, but it's had ITEM_LINKED, it must be active
                 * and have the key written to it already.
                 */
                hv = hash (ITEM_key (it), it->nkey);
                if ( ( hold_lock = item_trylock (hv) ) == NULL )
                {
                    status = MOVE_LOCKED;
                }
                else
                {
                    refcount = refcount_incr (&it->refcount);
                    if ( refcount == 2 )
                    { /* item is linked but not busy */
                        /* Double check ITEM_LINKED flag here, since we're
                         * past a memory barrier from the mutex. */
                        if ( ( it->it_flags & ITEM_LINKED ) != 0 )
                        {
                            status = MOVE_FROM_LRU;
                        }
                        else
                        {
                            /* refcount == 1 + !ITEM_LINKED means the item is being
                             * uploaded to, or was just unlinked but hasn't been freed
                             * yet. Let it bleed off on its own and try again later */
                            status = MOVE_BUSY;
                        }
                    }
                    else
                    {
                        if ( settings.verbose > 2 )
                        {
                            fprintf (stderr, "Slab reassign hit a busy item: refcount: %d (%d -> %d)\n",
                                     it->refcount, slab_rebal.s_clsid, slab_rebal.d_clsid);
                        }
                        status = MOVE_BUSY;
                    }
                    /* Item lock must be held while modifying refcount */
                    if ( status == MOVE_BUSY )
                    {
                        refcount_decr (&it->refcount);
                        item_trylock_unlock (hold_lock);
                    }
                }
            }
        }

        switch ( status )
        {
            case MOVE_FROM_LRU:
                /* Lock order is LRU locks -> slabs_lock. unlink uses LRU lock.
                 * We only need to hold the slabs_lock while initially looking
                 * at an item, and at this point we have an exclusive refcount
                 * (2) + the item is locked. Drop slabs lock, drop item to
                 * refcount 1 (just our own, then fall through and wipe it
                 */
                pthread_mutex_unlock (&slabs_lock);
                do_item_unlink (it, hv);
                item_trylock_unlock (hold_lock);
                pthread_mutex_lock (&slabs_lock);
            case MOVE_FROM_SLAB:
                it->refcount = 0;
                it->it_flags = 0;
                it->slabs_clsid = 255;
                break;
            case MOVE_BUSY:
            case MOVE_LOCKED:
                slab_rebal.busy_items ++;
                was_busy ++;
                break;
            case MOVE_PASS:
                break;
        }

        slab_rebal.slab_pos = ( char * ) slab_rebal.slab_pos + s_cls->size;
        if ( slab_rebal.slab_pos >= slab_rebal.slab_end )
            break;
    }

    if ( slab_rebal.slab_pos >= slab_rebal.slab_end )
    {
        /* Some items were busy, start again from the top */
        if ( slab_rebal.busy_items )
        {
            slab_rebal.slab_pos = slab_rebal.slab_start;
            slab_rebal.busy_items = 0;
        }
        else
        {
            slab_rebal.done ++;
        }
    }

    pthread_mutex_unlock (&slabs_lock);

    return was_busy;
}

static void slab_rebalance_finish ( void )
{
    slabclass_t *s_cls;
    slabclass_t *d_cls;

    pthread_mutex_lock (&slabs_lock);

    s_cls = & slabclass[slab_rebal.s_clsid];
    d_cls = & slabclass[slab_rebal.d_clsid];

    /* At this point the stolen slab is completely clear */
    s_cls->slab_list[s_cls->killing - 1] =
        s_cls->slab_list[s_cls->slabs - 1];
    s_cls->slabs --;
    s_cls->killing = 0;

    memset (slab_rebal.slab_start, 0, ( size_t ) settings.item_size_max);

    d_cls->slab_list[d_cls->slabs ++] = slab_rebal.slab_start;
    split_slab_page_into_freelist (slab_rebal.slab_start,
                                   slab_rebal.d_clsid);

    slab_rebal.done = 0;
    slab_rebal.s_clsid = 0;
    slab_rebal.d_clsid = 0;
    slab_rebal.slab_start = NULL;
    slab_rebal.slab_end = NULL;
    slab_rebal.slab_pos = NULL;

    slab_rebalance_signal = 0;

    pthread_mutex_unlock (&slabs_lock);

    if ( settings.verbose > 1 )
    {
        fprintf (stderr, "finished a slab move\n");
    }
}

/* Return 1 means a decision was reached.
 * Move to its own thread (created/destroyed as needed) once automover is more
 * complex.
 */
static int slab_automove_decision ( int *src, int *dst )
{
    static uint64_t evicted_old[MAX_NUMBER_OF_SLAB_CLASSES];
    static unsigned int slab_zeroes[MAX_NUMBER_OF_SLAB_CLASSES];
    static unsigned int slab_winner = 0;
    static unsigned int slab_wins = 0;
    uint64_t evicted_new[MAX_NUMBER_OF_SLAB_CLASSES];
    uint64_t evicted_diff = 0;
    uint64_t evicted_max = 0;
    unsigned int highest_slab = 0;
    unsigned int total_pages[MAX_NUMBER_OF_SLAB_CLASSES];
    int i;
    int source = 0;
    int dest = 0;
    static rel_time_t next_run;

    /* Run less frequently than the slabmove tester. */
    if ( current_time >= next_run )
    {
        next_run = current_time + 10;
    }
    else
    {
        return 0;
    }

    item_stats_evictions (evicted_new);
    pthread_mutex_lock (&slabs_lock);
    for ( i = POWER_SMALLEST; i < power_largest; i ++ )
    {
        total_pages[i] = slabclass[i].slabs;
    }
    pthread_mutex_unlock (&slabs_lock);

    /* Find a candidate source; something with zero evicts 3+ times */
    for ( i = POWER_SMALLEST; i < power_largest; i ++ )
    {
        evicted_diff = evicted_new[i] - evicted_old[i];
        if ( evicted_diff == 0 && total_pages[i] > 2 )
        {
            slab_zeroes[i] ++;
            if ( source == 0 && slab_zeroes[i] >= 3 )
                source = i;
        }
        else
        {
            slab_zeroes[i] = 0;
            if ( evicted_diff > evicted_max )
            {
                evicted_max = evicted_diff;
                highest_slab = i;
            }
        }
        evicted_old[i] = evicted_new[i];
    }

    /* Pick a valid destination */
    if ( slab_winner != 0 && slab_winner == highest_slab )
    {
        slab_wins ++;
        if ( slab_wins >= 3 )
            dest = slab_winner;
    }
    else
    {
        slab_wins = 1;
        slab_winner = highest_slab;
    }

    if ( source && dest )
    {
        *src = source;
        *dst = dest;
        return 1;
    }
    return 0;
}

/* Slab rebalancer thread.
 * Does not use spinlocks since it is not timing sensitive. Burn less CPU and
 * go to sleep if locks are contended
 */
static void *slab_maintenance_thread ( void *arg )
{
    int src, dest;

    while ( do_run_slab_thread )
    {
        if ( settings.slab_automove == 1 )
        {
            if ( slab_automove_decision (&src, &dest) == 1 )
            {
                /* Blind to the return codes. It will retry on its own */
                slabs_reassign (src, dest);
            }
            sleep (1);
        }
        else
        {
            /* Don't wake as often if we're not enabled.
             * This is lazier than setting up a condition right now. */
            sleep (5);
        }
    }
    return NULL;
}

/* Slab mover thread.
 * Sits waiting for a condition to jump off and shovel some memory about
 */
static void *slab_rebalance_thread ( void *arg )
{
    int was_busy = 0;
    /* So we first pass into cond_wait with the mutex held */
    mutex_lock (&slabs_rebalance_lock);

    while ( do_run_slab_rebalance_thread )
    {
        if ( slab_rebalance_signal == 1 )
        {
            if ( slab_rebalance_start () < 0 )
            {
                /* Handle errors with more specifity as required. */
                slab_rebalance_signal = 0;
            }

            was_busy = 0;
        }
        else if ( slab_rebalance_signal && slab_rebal.slab_start != NULL )
        {
            was_busy = slab_rebalance_move ();
        }

        if ( slab_rebal.done )
        {
            slab_rebalance_finish ();
        }
        else if ( was_busy )
        {
            /* Stuck waiting for some items to unlock, so slow down a bit
             * to give them a chance to free up */
            usleep (50);
        }

        if ( slab_rebalance_signal == 0 )
        {
            /* always hold this lock while we're running */
            pthread_cond_wait (&slab_rebalance_cond, &slabs_rebalance_lock);
        }
    }
    return NULL;
}

/* Iterate at most once through the slab classes and pick a "random" source.
 * I like this better than calling rand() since rand() is slow enough that we
 * can just check all of the classes once instead.
 */
static int slabs_reassign_pick_any ( int dst )
{
    static int cur = POWER_SMALLEST - 1;
    int tries = power_largest - POWER_SMALLEST + 1;
    for (; tries > 0; tries -- )
    {
        cur ++;
        if ( cur > power_largest )
            cur = POWER_SMALLEST;
        if ( cur == dst )
            continue;
        if ( slabclass[cur].slabs > 1 )
        {
            return cur;
        }
    }
    return - 1;
}

static enum reassign_result_type do_slabs_reassign ( int src, int dst )
{
    if ( slab_rebalance_signal != 0 )
        return REASSIGN_RUNNING;

    if ( src == dst )
        return REASSIGN_SRC_DST_SAME;

    /* Special indicator to choose ourselves. */
    if ( src == - 1 )
    {
        src = slabs_reassign_pick_any (dst);
        /* TODO: If we end up back at -1, return a new error type */
    }

    if ( src < POWER_SMALLEST || src > power_largest ||
         dst < POWER_SMALLEST || dst > power_largest )
        return REASSIGN_BADCLASS;

    if ( slabclass[src].slabs < 2 )
        return REASSIGN_NOSPARE;

    slab_rebal.s_clsid = src;
    slab_rebal.d_clsid = dst;

    slab_rebalance_signal = 1;
    pthread_cond_signal (&slab_rebalance_cond);

    return REASSIGN_OK;
}

enum reassign_result_type slabs_reassign ( int src, int dst )
{
    enum reassign_result_type ret;
    if ( pthread_mutex_trylock (&slabs_rebalance_lock) != 0 )
    {
        return REASSIGN_RUNNING;
    }
    ret = do_slabs_reassign (src, dst);
    pthread_mutex_unlock (&slabs_rebalance_lock);
    return ret;
}

/* If we hold this lock, rebalancer can't wake up or move */
void slabs_rebalancer_pause ( void )
{
    pthread_mutex_lock (&slabs_rebalance_lock);
}

void slabs_rebalancer_resume ( void )
{
    pthread_mutex_unlock (&slabs_rebalance_lock);
}

static pthread_t maintenance_tid;
static pthread_t rebalance_tid;

int start_slab_maintenance_thread ( void )
{
    int ret;
    slab_rebalance_signal = 0;
    slab_rebal.slab_start = NULL;
    char *env = getenv ("MEMCACHED_SLAB_BULK_CHECK");
    if ( env != NULL )
    {
        slab_bulk_check = atoi (env);
        if ( slab_bulk_check == 0 )
        {
            slab_bulk_check = DEFAULT_SLAB_BULK_CHECK;
        }
    }

    if ( pthread_cond_init (&slab_rebalance_cond, NULL) != 0 )
    {
        fprintf (stderr, "Can't intiialize rebalance condition\n");
        return - 1;
    }
    pthread_mutex_init (&slabs_rebalance_lock, NULL);

    if ( ( ret = pthread_create (&maintenance_tid, NULL,
                                 slab_maintenance_thread, NULL) ) != 0 )
    {
        fprintf (stderr, "Can't create slab maint thread: %s\n", strerror (ret));
        return - 1;
    }
    if ( ( ret = pthread_create (&rebalance_tid, NULL,
                                 slab_rebalance_thread, NULL) ) != 0 )
    {
        fprintf (stderr, "Can't create rebal thread: %s\n", strerror (ret));
        return - 1;
    }
    return 0;
}

/* The maintenance thread is on a sleep/loop cycle, so it should join after a
 * short wait */
void stop_slab_maintenance_thread ( void )
{
    mutex_lock (&slabs_rebalance_lock);
    do_run_slab_thread = 0;
    do_run_slab_rebalance_thread = 0;
    pthread_cond_signal (&slab_rebalance_cond);
    pthread_mutex_unlock (&slabs_rebalance_lock);

    /* Wait for the maintenance thread to stop */
    pthread_join (maintenance_tid, NULL);
    pthread_join (rebalance_tid, NULL);
}
